'''
预处理图像特征，对于每一张图片进行目标检测，将目标检测到的bbox使用resnet提取特征(2048维)，同标签一起制作成字典
{图片文件名：目标列表}
目标列表元素：字典：{'feature': ndarray(2048), 'label': int, 'possibility': float}
'''
import sys

import mmcv
import torch
import pickle
import os
from tqdm import tqdm
from torchvision import transforms
from model.encoders.MMDetection import MMDetection
from model.encoders.ResNet101 import ResNet101

TOP = 5
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transformation = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])


def crop_image(img, rect):
    img = mmcv.imcrop(img, rect[:4])
    return img


def image_resize_to_224(img):
    img = mmcv.imresize(img, (224, 224))
    return img


OD = MMDetection(DEVICE)
resnet = ResNet101(DEVICE)

for DATASET in ['twitter2015', 'twitter2017']:
    for TYPE in ['train', 'dev', 'test']:
        print(DATASET, TYPE)
        features = {}
        imgs = os.listdir(os.path.join('data', 'raw', DATASET, 'images', TYPE))

        for filename in tqdm(imgs, total=len(imgs), desc='Extracting', file=sys.stdout):
            img_path = os.path.join('data', 'raw', DATASET, 'images', TYPE, filename)
            img = mmcv.imread(img_path)
            V_o_labels = OD(img_path)
            objects = []
            for i in range(len(V_o_labels)):
                for num in range(len(V_o_labels[i])):
                    possibility = V_o_labels[i][num][4]
                    rect = V_o_labels[i][num][:4]
                    result = {
                        'rect': rect,
                        'label': i,
                        'possibility': possibility
                    }
                    objects.append(result)
            objects_sorted = sorted(objects, key=lambda x: x['possibility'], reverse=True)[:TOP]
            objects.clear()
            for bbox in objects_sorted:
                cropped_image = crop_image(img, bbox['rect'])
                cropped_image = image_resize_to_224(cropped_image)
                cropped_image = transformation(cropped_image)
                cropped_image = cropped_image.unsqueeze(0)
                cropped_image = cropped_image.to(DEVICE)
                with torch.no_grad():
                    feat, _ = resnet(cropped_image)
                feat = feat.detach().cpu().numpy()
                result = {
                    'feature': feat[0],
                    'label': bbox['label'],
                    'possibility': bbox['possibility']
                }
                objects.append(result)
            features[filename] = objects
        with open(os.path.join('dataset', DATASET, 'images', 'image_features_' + TYPE + '.pkl'), 'wb') as f:
            pickle.dump(features, f)
